{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-02-26T13:31:38.073766400Z",
     "start_time": "2024-02-26T13:31:38.050570900Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3 saveAsSVG\n",
      "4 File file = f;\n",
      "5 if (file == null) {\n",
      "37 if (file != null) {\n"
     ]
    }
   ],
   "source": [
    "import javalang\n",
    "\n",
    "java_code = r\"\"\"\n",
    "public class Test{\n",
    "    protected void saveAsSVG(File f) throws IOException {\n",
    "        File file = f;\n",
    "        if (file == null) {\n",
    "            JFileChooser fileChooser = new JFileChooser();\n",
    "            fileChooser.setCurrentDirectory(this.defaultDirectoryForSaveAs);\n",
    "            FileNameExtensionFilter filter = new FileNameExtensionFilter(\n",
    "                    localizationResources.getString(\"SVG_Files\"), \"svg\");\n",
    "            fileChooser.addChoosableFileFilter(filter);\n",
    "            fileChooser.setFileFilter(filter);\n",
    "\n",
    "            int option = fileChooser.showSaveDialog(this);\n",
    "            if (option == JFileChooser.APPROVE_OPTION) {\n",
    "                String filename = fileChooser.getSelectedFile().getPath();\n",
    "                if (isEnforceFileExtensions()) {\n",
    "                    if (!filename.endsWith(\".svg\")) {\n",
    "                        filename = filename + \".svg\";\n",
    "                    }\n",
    "                }\n",
    "                file = new File(filename);\n",
    "                if (file.exists()) {\n",
    "                    String fileExists = localizationResources.getString(\n",
    "                            \"FILE_EXISTS_CONFIRM_OVERWRITE\");\n",
    "                    int response = JOptionPane.showConfirmDialog(this,\n",
    "                            fileExists,\n",
    "                            localizationResources.getString(\"Save_as_SVG\"),\n",
    "                            JOptionPane.OK_CANCEL_OPTION);\n",
    "                    if (response == JOptionPane.CANCEL_OPTION) {\n",
    "                        file = null;\n",
    "                    }\n",
    "                }\n",
    "\n",
    "            }\n",
    "        }\n",
    "\n",
    "        if (file != null) {\n",
    "            // use reflection to get the SVG string\n",
    "            String svg = generateSVG(getWidth(), getHeight());\n",
    "            BufferedWriter writer = null;\n",
    "            Exception originalException = null;\n",
    "            try {\n",
    "                writer = new BufferedWriter(new FileWriter(file));\n",
    "                writer.write(\"<!DOCTYPE svg PUBLIC \\\"-//W3C//DTD SVG 1.1//EN\\\" \\\"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\\\">\\n\");\n",
    "                writer.write(svg + \"\\n\");\n",
    "                writer.flush();\n",
    "            } catch (Exception e) {\n",
    "                originalException = e;\n",
    "            }\n",
    "            try {\n",
    "                if (writer != null) {\n",
    "                    writer.close();\n",
    "                }\n",
    "            } catch (IOException ex) {\n",
    "                RuntimeException th = new RuntimeException(ex);\n",
    "                if (originalException != null)\n",
    "                    th.addSuppressed(originalException);\n",
    "                throw th;\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "\"\"\"\n",
    "lines = java_code.splitlines()\n",
    "tree = javalang.parse.parse(java_code)\n",
    "clazz = tree.types[0]\n",
    "method = clazz.methods[0]\n",
    "print(method.position[0], method.name)\n",
    "# 遍历method，找到所有语句的start\n",
    "for node in method.body:\n",
    "    start = node.position[0]\n",
    "    print(start, lines[start - 1].strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "data": {
      "text/plain": "{(3, '    protected void saveAsSVG(File f) throws IOException {'),\n (4, '        File file = f;'),\n (5, '        if (file == null) {'),\n (6, '            JFileChooser fileChooser = new JFileChooser();'),\n (7,\n  '            fileChooser.setCurrentDirectory(this.defaultDirectoryForSaveAs);'),\n (8,\n  '            FileNameExtensionFilter filter = new FileNameExtensionFilter('),\n (9,\n  '                    localizationResources.getString(\"SVG_Files\"), \"svg\");'),\n (10, '            fileChooser.addChoosableFileFilter(filter);'),\n (11, '            fileChooser.setFileFilter(filter);'),\n (13, '            int option = fileChooser.showSaveDialog(this);'),\n (14, '            if (option == JFileChooser.APPROVE_OPTION) {'),\n (15,\n  '                String filename = fileChooser.getSelectedFile().getPath();'),\n (16, '                if (isEnforceFileExtensions()) {'),\n (17, '                    if (!filename.endsWith(\".svg\")) {'),\n (18, '                        filename = filename + \".svg\";'),\n (21, '                file = new File(filename);'),\n (22, '                if (file.exists()) {'),\n (23,\n  '                    String fileExists = localizationResources.getString('),\n (24, '                            \"FILE_EXISTS_CONFIRM_OVERWRITE\");'),\n (25,\n  '                    int response = JOptionPane.showConfirmDialog(this,'),\n (26, '                            fileExists,'),\n (27,\n  '                            localizationResources.getString(\"Save_as_SVG\"),'),\n (28, '                            JOptionPane.OK_CANCEL_OPTION);'),\n (29, '                    if (response == JOptionPane.CANCEL_OPTION) {'),\n (30, '                        file = null;'),\n (37, '        if (file != null) {'),\n (39, '            String svg = generateSVG(getWidth(), getHeight());'),\n (40, '            BufferedWriter writer = null;'),\n (41, '            Exception originalException = null;'),\n (42, '            try {'),\n (43, '                writer = new BufferedWriter(new FileWriter(file));'),\n (44,\n  '                writer.write(\"<!DOCTYPE svg PUBLIC \\\\\"-//W3C//DTD SVG 1.1//EN\\\\\" \\\\\"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\\\\\">\\\\n\");'),\n (45, '                writer.write(svg + \"\\\\n\");'),\n (46, '                writer.flush();'),\n (48, '                originalException = e;'),\n (50, '            try {'),\n (51, '                if (writer != null) {'),\n (52, '                    writer.close();'),\n (55, '                RuntimeException th = new RuntimeException(ex);'),\n (56, '                if (originalException != null)'),\n (57, '                    th.addSuppressed(originalException);'),\n (58, '                throw th;')}"
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# iter method\n",
    "def get_method_end(method):\n",
    "    def f(node):\n",
    "        if hasattr(node, 'children'):\n",
    "            for child in node.children:\n",
    "                yield from f(child)\n",
    "\n",
    "        if isinstance(node, list):\n",
    "            for child in node:\n",
    "                yield from f(child)\n",
    "\n",
    "        if hasattr(node, 'position'):\n",
    "            if node.position is not None:\n",
    "                try:\n",
    "                    yield node.position[0], lines[node.position[0] - 1]\n",
    "                except Exception as e:\n",
    "                    print(node.position)\n",
    "    stmt_set = set(f(method))\n",
    "    return max(stmt_set, key=lambda x: x[0]) + 1"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-26T13:35:46.014004900Z",
     "start_time": "2024-02-26T13:35:46.009685700Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "outputs": [
    {
     "data": {
      "text/plain": "['import org.jfree.chart.plot.Plot;']"
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "def get_class_imports(source_folder, class_names: list):\n",
    "    # 递归遍历source文件夹，找到所有的class文件\n",
    "    class_map = {} # {class_name: [package_reference, ...]}\n",
    "    for root, dirs, files in os.walk(source_folder):\n",
    "        for file in files:\n",
    "            if file.endswith(\".java\"):\n",
    "                class_name = file[:-5]\n",
    "                root = root.replace(\"\\\\\", \".\")\n",
    "                root = root.replace(\"/\", \".\")\n",
    "                if \"src.test.java\" in root:\n",
    "                    continue\n",
    "                ref = root.split(\"src.main.java.\")[1]\n",
    "                if class_name not in class_map:\n",
    "                    class_map[class_name] = [ref]\n",
    "                else:\n",
    "                    class_map[class_name].append(ref)\n",
    "    imports = []\n",
    "    for name in class_names:\n",
    "        if name not in class_map:\n",
    "            return None\n",
    "\n",
    "        refs = class_map[name]\n",
    "        for ref in refs:\n",
    "            imports.append(f\"import {ref}.{name};\")\n",
    "    return imports\n",
    "get_class_imports(r\"D:\\IDEA_Projects\\TestJavaCode\\src\", [\"Plot\"])"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-02-28T03:54:05.781446500Z",
     "start_time": "2024-02-28T03:54:05.769808200Z"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
